__doc__ =   """ Class to create PDB files from entities.
            """

from collections import defaultdict
from itertools import chain

from Einput import einput
from ResidueProperties import AA_NAMES
from PDBParser import dict2pdb, dict2ter


def write_header(header):
    if not isinstance(header, dict):
        return (header or [])
    xt = ('REMARK 200 EXPERIMENT TYPE : %s\n', 'experiment_type')
    sg = ('REMARK 290 SYMMETRY OPERATORS FOR SPACE GROUP: %s\n', 'space_group')
    tests = [xt, sg]
    results = []
    for test in tests:
        string, key = test
        try:
            value = header[key]
            results.append(string % value)
        except (KeyError, TypeError):
            pass
    return results

def write_coords(atoms):
    old_fields = defaultdict(str)
    coords = ['MODEL        1\n']
    for atom in atoms.sorted_values():
        fields = atom.get_dict()
        if (old_fields['chain_id'] != fields['chain_id']) and old_fields['chain_id'] and old_fields['res_name'] in AA_NAMES:
                coords.append(dict2ter(old_fields))
        if (old_fields['model'] != fields['model']) and old_fields['model'] != '': # model can be 0 :)
                if old_fields['chain_id'] and old_fields['res_name'] in AA_NAMES:
                    coords.append(dict2ter(old_fields))
                coords.append('ENDMDL\n')
                coords.append('MODEL     %4i\n' % (fields['model'] + 1))
        coords.append(dict2pdb(fields))
        old_fields = fields
    if fields['res_name'] in AA_NAMES:
        coords.append(dict2ter(fields))
    coords.append('ENDMDL\n')
    return coords

def write_trailer(trailer):
    if not isinstance(trailer, dict):
        return (trailer or [])


class PDBWriter(object):

    def set_structure(self, entities, header_ =None, trailer_ =None):
        self.structure = einput(entities, level ='A', name ='atoms')
        # hierarchy: args, dicts
        try:
            header = (header_ or entities.raw_header or entities.header)
        except AttributeError:
            header = header_
        try:
            trailer = (trailer_ or entities.raw_trailer or entities.trailer)
        except AttributeError:
            trailer = trailer_

        self.coords = write_coords(self.structure)
        self.header = write_header(header)
        self.trailer = write_trailer(trailer)

    def save(self, file):
        self.file = open(file, 'w')
        for part in chain(self.header, self.coords, self.trailer):
            self.file.writelines(part)
        self.file.close()


#EOF